# _*_ coding: utf-8 _*_

import numpy as np
from copy import deepcopy
from random import uniform

import matplotlib.pyplot as plt

from scripts.utils.operators import Lifecycle
from scripts.utils.operators import lhs_uniform

from data.CEC.CEC2017.cec17_functions import cec17_test_func

np.set_printoptions(linewidth=999, threshold=1e9)


class PSO:
    """
    Particle Swarm Optimization (PSO)
    """

    def __init__(self, obj_func, dimension: int, bounds: tuple, NP: int, w: float, c1: float, c2: float):
        self.obj_func = obj_func
        self.__d, self.__bounds = dimension, bounds
        self.__NP, self.__w, self.__c1, self.__c2 = NP, w, c1, c2

        lhs = lhs_uniform(n=self.__NP, p=self.__d)
        x = (self.__bounds[0] + lhs * (self.__bounds[1] - self.__bounds[0]))  # position
        v = np.zeros(shape=(self.__NP, self.__d))  # velocity
        # population[i] structure: [x, v, fitness]
        self.__population = [[x[i], v[i], 0] for i in range(self.__NP)]
        self.__evaluate()
        self.pBest = deepcopy(min(self.__population, key=lambda _: _[2]))
        self.gBest = deepcopy(self.pBest)

        self.life = Lifecycle()
        self.life.update(eva=self.__NP)

    def optimize(self, iterations=None, evaluations=None, deadline=None):

        self.life.set_stopping_condition(iterations=iterations, evaluations=evaluations, deadline=deadline)
        while self.life.check_stopping_condition() is False:

            self.__learning()
            self.__evaluate()

            self.pBest = deepcopy(min(self.__population, key=lambda x: x[2]))

            if self.gBest[2] > self.pBest[2]:
                self.gBest = deepcopy(self.pBest)
            self.life.update(it=1, eva=self.__NP, gBest=self.gBest[2])

        return self.life

    def __evaluate(self):
        for i, x in enumerate(self.__population):
            self.__population[i][2] = self.obj_func(x=self.__population[i][0])

    def __learning(self):
        for i in range(self.__NP):
            x1, x2 = self.pBest[0], self.gBest[0]
            self.__population[i][1], self.__population[i][0] = self.__update(self.__population[i], x1, x2)

    def __update(self, p, x1, x2):
        """
        update velocity and position
        :param p: current particle
        :param x1: pBest position
        :param x2: gBest position
        :return: velocity and position
        """
        x, v = p[0], p[1]
        r1 = np.random.uniform(low=0.0, high=1.0, size=self.__d)
        r2 = np.random.uniform(low=0.0, high=1.0, size=self.__d)
        v = self.__w * v + self.__c1 * r1 * (x1 - x) + self.__c2 * r2 * (x2 - x)
        x = x + v
        for i in range(len(x)):
            if x[i] < self.__bounds[0]:
                x[i] = self.__bounds[0]
            if x[i] > self.__bounds[1]:
                x[i] = self.__bounds[1]
        return v, x
